import torch
from torch import nn
import torch.nn.init as init


class ReGru(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layer):
        super(ReGru, self).__init__()
        self.num_layer = num_layer
        self.hidden_size = hidden_size

        self.input_transform = nn.Linear(input_size, hidden_size)
        self.gru_layers = nn.ModuleList([
            nn.GRU(hidden_size, hidden_size, batch_first=True) for _ in range(num_layer)
        ])

        self.norm_layers = nn.ModuleList([
            nn.LayerNorm(hidden_size) for _ in range(num_layer)
        ])
        
        self.fc = nn.Linear(hidden_size, output_size)
        

    def forward(self, x):
        residual = self.input_transform(x)

        for i in range(self.num_layer):
            out, _ = self.gru_layers[i](residual)
            residual = self.norm_layers[i](residual + out)

        # 只使用序列的最后一个输出
        out = self.fc(residual[:, -1, :])
        return out
    
    def predict(self, x):
        h0 = torch.zeros(self.num_layer, self.hidden_size).to(x.device)
        out, _ = self.gru(x, h0)
        print("gru:")
        print(out.shape)
        print(out)
        print(_.shape)
        print(_)
        print("gru wieght")
        print(self.gru)
        print("gru end")
        out = self.fc(out[-1, :])
        return out
